package avicit.bdp.dds.common.utils;

import avicit.bdp.common.dto.CalculateEngineConf;
import avicit.bdp.common.utils.ConfigUtils;
import avicit.bdp.core.http.HttpClientUtil;
import avicit.bdp.core.http.common.HttpConfig;
import avicit.bdp.core.http.exception.HttpProcessException;
import avicit.bdp.core.util.json.JSONUtils;
import avicit.bdp.dds.common.Constants;
import avicit.bdp.dds.common.utils.webhdfs.KerberosWebYarnUtils;
import avicit.bdp.dds.dispatch.enums.ExecutionStatus;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.yarn.client.cli.RMAdminCLI;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.util.HashMap;
import java.util.Map;

/**
 * TODO
 *
 * @author xugb
 * @date 2021/6/10 15:31
 */
public class YarnUtils {
    private static final Logger logger = LoggerFactory.getLogger(YarnUtils.class);
    public static final String rmHaIds = ConfigUtils.getInstance().getString(avicit.bdp.core.constant.Constants.YARN_RESOURCEMANAGER_HA_RM_IDS);
    public static final String APP_ADDRESS = ConfigUtils.getInstance().getString(avicit.bdp.core.constant.Constants.YARN_APPLICATION_STATUS_ADDRESS);
    private static volatile boolean yarnEnabled = false;

    private static final String YARN_ADDRESS_TEMPLATE = "http://%s:%d/ws/v1/cluster/apps/%s";

    public static String killYarnApp(CalculateEngineConf conf, String applicationId) throws Exception {
        if (StringUtils.isEmpty(applicationId)) {
            return null;
        }
        String appAddress = conf.getYarnAddress();

        boolean isKerberos = Constants.KERBEROS.equalsIgnoreCase(conf.getAuthType());

        String responseContent = "";
        String applicationUrl = "";
        if (isKerberos) {
            KerberosWebYarnUtils utils = new KerberosWebYarnUtils(appAddress, conf.getKrb5Path(), conf.getKeytabPath(), conf.getKerberosUser());
            responseContent = utils.killApp(applicationId);
        } else {

            if (appAddress.endsWith(Constants.SINGLE_SLASH)) {
                applicationUrl = appAddress + "ws/v1/cluster/apps/" + applicationId + "/state";
            } else {
                applicationUrl = appAddress + "/ws/v1/cluster/apps/" + applicationId + "/state";
            }

            Map<String, Object> map = new HashMap<>(2);
            map.put("state", "KILLED");

            HttpConfig config = HttpConfig.custom().timeout(1000).url(applicationUrl).json(JSONUtils.toJson(map));

            responseContent = HttpClientUtil.put(config);
        }
        logger.info(applicationUrl + " : " + responseContent);
        return responseContent;
    }

    private static String getYarnAddress(String url, String applicationId) {
        if (!isurl(url)) {
            throw new IllegalArgumentException("url参数不合法!");
        }
        URI uri = URI.create(url);
        String host = uri.getHost();
        int port = uri.getPort();
        return String.format(YARN_ADDRESS_TEMPLATE, host, port, applicationId);
    }

    public static boolean isurl(String str) {
        str = str.toLowerCase();
        String regex = "^((https|http|ftp|rtsp|mms)?://)"
                + "?(([0-9a-z_!~*'().&=+$%-]+: )?[0-9a-z_!~*'().&=+$%-]+@)?"
                + "(([0-9]{1,3}\\.){3}[0-9]{1,3}"
                + "|"
                + "([0-9a-z_!~*'()-]+\\.)*"
                + "([0-9a-z][0-9a-z-]{0,61})?[0-9a-z]\\."
                + "[a-z]{2,6})"
                + "(:[0-9]{1,5})?"
                + "((/?)|"
                + "(/[0-9a-z_!~*'().;?:@&=+$,%#-]+)+/?)$";
        return str.matches(regex);
    }


    public static ExecutionStatus getApplicationStatus(CalculateEngineConf conf, String applicationId) throws Exception {
        String applicationUrl = getApplicationUrl(conf, applicationId);
        if (applicationUrl == null) {
            return null;
        }

        boolean isKerberos = Constants.KERBEROS.equalsIgnoreCase(conf.getAuthType());
        String responseContent = "";
        if (isKerberos) {
            HttpKerberosUtils utils = new HttpKerberosUtils(conf.getKerberosUser(), conf.getKeytabPath(), conf.getKrb5Path(), false);
            responseContent = utils.get(applicationUrl);
        } else {
            HttpConfig config = HttpConfig.custom().timeout(1000).url(applicationUrl);
            responseContent = HttpClientUtil.get(config);
        }

        JSONObject jsonObject = JSON.parseObject(responseContent);
        String result = jsonObject.getJSONObject("app").getString("finalStatus");

        switch (result) {
            case Constants.ACCEPTED:
                return ExecutionStatus.SUBMITTED_SUCCESS;
            case Constants.SUCCEEDED:
                return ExecutionStatus.SUCCESS;
            case Constants.NEW:
            case Constants.NEW_SAVING:
            case Constants.SUBMITTED:
            case Constants.FAILED:
                return ExecutionStatus.FAILURE;
            case Constants.KILLED:
                return ExecutionStatus.KILL;

            case Constants.RUNNING:
            default:
                return ExecutionStatus.RUNNING_EXEUTION;
        }
    }

    private static String getApplicationUrl(CalculateEngineConf conf, String applicationId) {
        if (StringUtils.isEmpty(applicationId)) {
            return null;
        }

        String yarnAddress = conf.getYarnAddress();

        if (StringUtils.isEmpty(yarnAddress)) {
            return null;
        }

        String applicationUrl = getYarnAddress(yarnAddress, applicationId);
        logger.info("applicationUrl={}", applicationUrl);
        return applicationUrl;
    }

    /**
     * get application url
     *
     * @param applicationId application id
     * @return url of application
     */
    public static String getApplicationUrl(String applicationId) throws Exception {
        /*
         * if rmHaIds contains xx, it signs not use resourcemanager
         * otherwise:
         *  if rmHaIds is empty, single resourcemanager enabled
         *  if rmHaIds not empty: resourcemanager HA enabled
         */
        String appUrl = "";

        if (StringUtils.isEmpty(rmHaIds)) {
            //single resourcemanager enabled
            appUrl = APP_ADDRESS;
            yarnEnabled = true;
        } else {
            //resourcemanager HA enabled
            appUrl = getAppAddress(APP_ADDRESS, rmHaIds);
            yarnEnabled = true;
            logger.info("application url : {}", appUrl);
        }

        if (StringUtils.isBlank(appUrl)) {
            throw new Exception("application url is blank");
        }
        return String.format(appUrl, applicationId);
    }


    /**
     * getAppAddress
     *
     * @param appAddress app address
     * @param rmHa       resource manager ha
     * @return app address
     */
    public static String getAppAddress(String appAddress, String rmHa) {

        //get active ResourceManager
        String activeRM = YarnHAAdminUtils.getAcitveRMName(rmHa);

        String[] split1 = appAddress.split(Constants.DOUBLE_SLASH);

        if (split1.length != 2) {
            return null;
        }

        String start = split1[0] + Constants.DOUBLE_SLASH;
        String[] split2 = split1[1].split(Constants.COLON);

        if (split2.length != 2) {
            return null;
        }

        String end = Constants.COLON + split2[1];

        return start + activeRM + end;
    }

    /**
     * yarn ha admin utils
     */
    private static final class YarnHAAdminUtils extends RMAdminCLI {

        /**
         * get active resourcemanager
         *
         * @param rmIds
         * @return
         */
        public static String getAcitveRMName(String rmIds) {

            String[] rmIdArr = rmIds.split(Constants.COMMA);

            int activeResourceManagerPort = ConfigUtils.getInstance().getInteger(Constants.HADOOP_RESOURCE_MANAGER_HTTPADDRESS_PORT, 8088);

            String yarnUrl = "http://%s:" + activeResourceManagerPort + "/ws/v1/cluster/info";

            String state = null;
            try {
                /**
                 * send http get request to rm1
                 */
                state = getRMState(String.format(yarnUrl, rmIdArr[0]));

                if (Constants.HADOOP_RM_STATE_ACTIVE.equals(state)) {
                    return rmIdArr[0];
                } else if (Constants.HADOOP_RM_STATE_STANDBY.equals(state)) {
                    state = getRMState(String.format(yarnUrl, rmIdArr[1]));
                    if (Constants.HADOOP_RM_STATE_ACTIVE.equals(state)) {
                        return rmIdArr[1];
                    }
                } else {
                    return null;
                }
            } catch (Exception e) {
                state = getRMState(String.format(yarnUrl, rmIdArr[1]));
                if (Constants.HADOOP_RM_STATE_ACTIVE.equals(state)) {
                    return rmIdArr[0];
                }
            }
            return null;
        }


        /**
         * get ResourceManager state
         *
         * @param url
         * @return
         */
        public static String getRMState(String url) {
            String retStr = null;
            HttpConfig config = HttpConfig.custom().timeout(1000).url(url);
            try {
                retStr = HttpClientUtil.get(config);
            } catch (HttpProcessException e) {
                e.printStackTrace();
            }

            if (StringUtils.isEmpty(retStr)) {
                return null;
            }
            //to json
            JSONObject jsonObject = JSON.parseObject(retStr);

            //get ResourceManager state
            return jsonObject.getJSONObject("clusterInfo").getString("haState");
        }

    }
}
